#include <mlgesture.h>
#include <svm.h>
#include <hmm.h>
#include <trajectory.h>
#include <iostream>
#include <string>

#include <QFile>
#include <QString>
#include <QStringList>
#include <QTextStream>
#include <QMap>
#include <ctime>
#include <cstdlib>
#include <algorithm>

#include "../handMove/xmltrajectoryreader.h"

#include <boost/lexical_cast.hpp>

using boost::lexical_cast;
using boost::bad_lexical_cast;
using namespace std;
using namespace mlgesture;

enum GarbageType {
	LINEAR_GARBAGE = 0,
	FULL_GARBAGE = 1
};

void testStatic(const char* data_file, const char* label_file, double C, double gamma);
float staticCross(const vector<vector<double> > &train, const vector<vector<double> > &test, vector<vector<int> > &result,
				  const QMap<int, int> &labels, float &without_prob, double C, double gamma);
void testDynamic(const char* data_file, int states_count, GarbageType garbage);
float dynamicCross(const vector<Trajectory> &train, const vector<Trajectory> &test, vector<vector<int> > &result,
				   const vector<string> &types, float &without_prob, int state_count, GarbageType garbage);

void usage(const char* program) {
	cout << "Użycie: \n" << endl;
	cout << "Gesty statyczne: " << endl;
	cout << "\t" << program << " static static_data_file static_label_file C gamma" << endl;
	cout << "Gesty dynamiczne: " << endl;
	cout << "\t" << program << " dynamic dynamic_data_file states_cout garbage\n" << endl;
	cout << "static_data_file - plik z danymi statycznymi" << endl;
	cout << "static_label_file - plik z nazwamy etykiet" << endl;
	cout << "C - parametr klasyfikatora SVM - koszt błędnej klasyfikacji" << endl;
	cout << "gamma - parametr klasyfikatora SVM - współczynnik funkcji jądra" << endl;
	cout << "dynamic_data_file - plik xml danymi trajektorii" << endl;
	cout << "states_cout - ilość stanów w modelu" << endl;
	cout << "garbage - rodzaj modelu ogólnego; dopuszczalne wartości: linear - model prawo-lewy; full - model pełny" << endl;
}

int main(int argc, char *argv[]) {
	if (argc < 2) {
		cout << "Za mało parametrów" << endl;
		usage(argv[0]);
		return 0;
	}
	int type = -1;
	string type_s = lexical_cast<string>(argv[1]);
	if (type_s.compare("static")==0)
		type = 0;
	else if (type_s.compare("dynamic")==0)
		type = 1;
	else {
		usage(argv[0]);
		return 0;
	}
	srand(time(0));

	if (type == 0) {
		if (argc < 6) {
			usage(argv[0]);
			return 0;
		}

		double C=0, gamma=0;
		try {
			C = lexical_cast<double>(argv[4]);
			gamma = lexical_cast<double>(argv[5]);
		}
		catch(bad_lexical_cast&){
			cout << "Błędne parametry\n";
			usage(argv[0]);
			return 0;
		}

		testStatic(argv[2], argv[3], C, gamma);
	}
	else if (type == 1) {
		if (argc < 5) {
			usage(argv[0]);
			return 0;
		}

		int state_count;
		GarbageType garbage;
		string garbage_s;
		try {
			state_count = lexical_cast<int>(argv[3]);
			garbage_s = lexical_cast<string>(argv[4]);
		}
		catch(bad_lexical_cast&){
			cout << "Błędne parametry\n";
			usage(argv[0]);
			return 0;
		}
		if (garbage_s.compare("linear")==0)
			garbage = LINEAR_GARBAGE;
		else if (garbage_s.compare("full")==0)
			garbage = FULL_GARBAGE;
		else {
			usage(argv[0]);
			return 0;
		}
		testDynamic(argv[2], state_count, garbage);
	}
	return 0;
}

void testStatic(const char* data_file_c, const char* label_file_c, double C, double gamma) {
	QString data_file(data_file_c);
	QString labels_file(label_file_c);

	QMap<int, QString> labels;
	QMap<int, int> labels_id;
	QMap<int, int> groupCount;

	QFile data_qfile(labels_file);
	data_qfile.open(QIODevice::ReadOnly | QIODevice::Text);
	if (!data_qfile.isOpen()) {
		cout << "Nie można otworzyć pliku etykiet: " << labels_file.toStdString() << endl;
		return;
	}

	QTextStream in_data(&data_qfile);
	QString line = in_data.readLine();

	int label_tmp = 0;
	while (!line.isNull()) {
		QStringList lists = line.split(" ");
		if (lists.size() != 2) {
			cout << "Błędna struktura pliku etykiet" << endl;
			return;
		}
		bool ok;
		int tmp_id = lists[1].toInt(&ok);
		if (!ok) {
			cout << "Błędna struktura pliku etykiet" << endl;
			return;
		}
		if (!labels.contains(tmp_id)) {
			labels[tmp_id] = lists[0];
			labels_id[tmp_id] = label_tmp;
			groupCount[tmp_id] = 0;
			label_tmp++;
		}
		line = in_data.readLine();
	}
	data_qfile.close();

	data_qfile.setFileName(data_file);
	data_qfile.open(QIODevice::ReadOnly | QIODevice::Text);
	if (!data_qfile.isOpen()) {
		cout << "Nie można otworzyć pliku z danymi: " << data_file.toStdString() << endl;
		return;
	}

	in_data.setDevice(&data_qfile);

	int count=0;
	vector<vector<double> > data;
	line = in_data.readLine();
	while (!line.isNull()) {
		QStringList list = line.split(" ");
		if (list.size() != 8) {
			cout << "Błędna struktura pliku danych" << endl;
			return;
		}
		count++;
		line = in_data.readLine();
	}
	data.resize(count);
	in_data.seek(0);
	line = in_data.readLine();
	int line_num = 0;
	while (!line.isNull()) {
		QStringList list = line.split(" ");
		data[line_num].push_back(list[0].toInt());
		groupCount[data[line_num][0]]++;
		for (int i=0; i<7; i++) {
			QStringList list2 = list[i+1].split(":");
			if (list2.size() != 2) {
				cout << "Błędna struktura pliku danych" << endl;
				return;
			}
			data[line_num].push_back(list2[1].toDouble());
		}
		line_num++;
		line = in_data.readLine();
	}
	data_qfile.close();

	vector<vector<int> > res(labels.size());
	for(int i=0; i<labels.size(); i++) {
		res[i].resize(labels.size()+1);
		for(int j=0; j<=labels.size(); j++) {
			res[i][j]=0;
		}
	}
	vector<vector<double> > train, tests;

	random_shuffle(data.begin(), data.end());

	int data_size = data.size();
	float success_prob = 0;
	float without_prob = 0;
	float error_prob = 0;
	for (int i=0; i<30; i++) {
		train.clear();
		tests.clear();
		int start_test = i*data.size()/30  + (i < data_size%30 ? i : data.size()%30);
		int end_test = start_test + data_size/30 + (i < data_size%30 ? 1 : 0);

		int j = 0;
		for (vector<vector<double> >::iterator it = data.begin(); it != data.end(); ++it, ++j) {
			if (j>=start_test && j<end_test)
				tests.push_back(*it);
			else
				train.push_back(*it);
		}

		float tmp_prob1, tmp_prob2;
		tmp_prob1= staticCross(train, tests, res, labels_id, tmp_prob2, C, gamma);
		success_prob += tmp_prob1;
		without_prob += tmp_prob2;
		error_prob += (1-tmp_prob1-tmp_prob2);

	}
	cout << "Correct: " << success_prob/30 << endl;
	cout << "Error: " << error_prob/30 << endl;
	cout << "Bez decyzji: " << without_prob/30 << endl<<endl;


	cout << "Wyniki: " << endl;
	for (int i=0; i<labels.size(); i++) {
		for (int j=0; j<labels.size()+1; j++) {
			cout << res[i][j] << " ";
		}
		cout << endl;
	}

	cout << "\nStatystyka:" << endl;
	for (QMap<int, int>::iterator it = groupCount.begin(); it != groupCount.end(); ++it) {
		int key = it.key();
		cout << labels[key].toStdString() << ": " << it.value() << endl;
	}
}

float staticCross(const vector<vector<double> > &train, const vector<vector<double> > &test, vector<vector<int> > &result,
				   const QMap<int,int> &labels, float &without_prob, double C, double gamma) {
	svm_parameter param;
	svm_problem prob;
	svm_model *model;

	//set libsvmparams
	param.svm_type = C_SVC;
	param.kernel_type = RBF;
	param.degree = 3;
	param.gamma = gamma;
	param.coef0 = 0;
	param.nu = 0.5;
	param.cache_size = 100;
	param.C = C;
	param.eps = 1e-3;
	param.p = 0.1;
	param.shrinking = 1;
	param.probability = 1;
	param.nr_weight = 0;
	param.weight_label = NULL;
	param.weight = NULL;
	prob.l = train.size();

	prob.y = new double[prob.l];
	prob.x = new svm_node*[prob.l];

	for (int i=0; i<train.size(); i++) {
		prob.y[i] = train[i][0];
		prob.x[i] = new svm_node[8];
		for (int j=0; j<7; j++) {
			prob.x[i][j].index = j+1;
			prob.x[i][j].value = train[i][j+1];
		}
		prob.x[i][7].index = -1;
	}

	model = svm_train(&prob, &param);
	int* svm_labels = new int[result.size()];
	svm_get_labels(model, svm_labels);
	double* pred_prob = new double[result.size()];
	svm_node sample[8];
	int success=0, without=0;

	for (int i=0; i<test.size(); i++) {
		for (int j=0; j<8; j++) {
			sample[j].index = j+1;
			sample[j].value = test[i][j+1];
		}
		sample[7].index = -1;

		int res = svm_predict_probability(model, sample, pred_prob);
		double wyn_prob=0;
		for (int j=0; j<result.size(); j++) {
			if (svm_labels[j] == res) {
				wyn_prob = pred_prob[j];
				break;
			}
		}
		if (wyn_prob <= 0.5) {
			without++;
			result[labels[test[i][0]]][result.size()]++;
		}
		else if (res == test[i][0]) {
			success++;
			result[labels[test[i][0]]][labels[res]]++;
		}
		else {
			result[labels[test[i][0]]][labels[res]]++;
		}
	}
	without_prob = 1.*without/test.size();

	svm_free_model_content(model);

	delete [] svm_labels;
	delete [] pred_prob;
	for (int i=0; i<prob.l; i++) {
		delete [] prob.x[i];
	}
	delete [] prob.x;
	delete [] prob.y;
	return 1.*success/test.size();
}

void testDynamic(const char *data_file, int states_count, GarbageType garbage) {
	QFile file(data_file);
	if (!file.open(QFile::ReadOnly | QFile::Text)) {
		cout << "Nie można otworzyć pliku z danymi gestów dynamicznych" << endl;
		return;
	}

	vector<Trajectory> tracks;
	XmlTrajectoryReader reader(&tracks);
	if (!reader.read(&file)) {
		cout << "Zła struktura pliku z danymi gestów dynamicznych" << endl;
		return;
	}
	for(vector<Trajectory>::iterator it=tracks.begin(); it!=tracks.end();) {
		if (it->type_ == "garbage" || it->type_ == "undefined") {
			it = tracks.erase(it);
		}
		else ++it;
	}

	vector<string> labels;
	Trajectory::getTrajectoryTypes(tracks,labels);

	random_shuffle(tracks.begin(), tracks.end());

	QMap<string, int> groups_count;
	for(vector<Trajectory>::iterator it = tracks.begin(); it != tracks.end(); ++it) {
		if(!groups_count.contains(it->type_)) {
			groups_count[it->type_] = 0;
		}
		groups_count[it->type_]++;
	}

	vector<vector<int> > res(labels.size());
	for(int i=0; i<labels.size(); i++) {
		res[i].resize(labels.size()+1);
		for(int j=0; j<=labels.size(); j++) {
			res[i][j]=0;
		}
	}
	vector<Trajectory> train, tests;
	int data_size = tracks.size();
	float success_prob = 0;
	float without_prob = 0;
	float error_prob = 0;
	for (int i=0; i<20; i++) {
		train.clear();
		tests.clear();
		int start_test = i*data_size/20  + (i < data_size%20 ? i : data_size%20);
		int end_test = start_test + data_size/20 + (i < data_size%20 ? 1 : 0);

		int j = 0;
		for (vector<Trajectory>::iterator it = tracks.begin(); it != tracks.end(); ++it, ++j) {
			if (j>=start_test && j<end_test)
				tests.push_back(*it);
			else
				train.push_back(*it);
		}

		float tmp_prob1, tmp_prob2;
		tmp_prob1 = dynamicCross(train, tests, res, labels, tmp_prob2, states_count, garbage);
		success_prob += tmp_prob1;
		without_prob += tmp_prob2;
		error_prob += (1-tmp_prob1-tmp_prob2);
		cout  << "Grupa " << i << ":" << tmp_prob1 << " " << tmp_prob2 << " " << (1-tmp_prob1-tmp_prob2) << endl;
	}
	cout << "Correct: " << success_prob/20 << endl;
	cout << "Error: " << error_prob/20 << endl;
	cout << "Bez decyzji: " << without_prob/20 << endl<<endl;


	cout << "Wyniki: " << endl;
	for (int i=0; i<labels.size(); i++) {
		cout << labels[i] << ": ";
		for (int j=0; j<labels.size()+1; j++) {
			cout << res[i][j] << " ";
		}
		cout << endl;
	}

	cout << "\nStatystyka:" << endl;
	for (QMap<string, int>::iterator it = groups_count.begin(); it != groups_count.end(); ++it) {
		cout << it.key() << ": " << it.value() << endl;
	}
}

float dynamicCross(const vector<Trajectory> &train, const vector<Trajectory> &test, vector<vector<int> > &result,
				  const vector<string> &types, float &without_prob, int state_count, GarbageType garbage) {

	vector<vector<int> > data_train, garbage_train;
	vector<Hmm*> models;

	QMap<string, int> types_s;
	for (int i=0; i<types.size(); i++) {
		types_s[types[i]] = i;
	}

	for (size_t k=0; k<types.size(); k++) {
		data_train.clear();
		if (types[k] == "garbage" || types[k] == "undefined")
			continue;

		int count = 0;
		for (size_t i=0; i<train.size(); i++) {
			if (QString::fromStdString(train[i].type_).toLower() == QString::fromStdString(types[k]).toLower()) {
				count++;
				garbage_train.push_back(train[i].data_);
				data_train.push_back(train[i].data_);
			}
		}
		models.push_back(new Hmm(state_count, 16));
		models[k]->baumWelsh(data_train);
	}
	Hmm* garbage_model = new Hmm(state_count,16, garbage==LINEAR_GARBAGE ? Hmm::LINEAR_MODEL : Hmm::FULL_CONNECTED_MODEL);
	garbage_model->baumWelsh(garbage_train);

	int success=0, without=0;
	for (int i=0; i<test.size(); i++) {
		double max_prob = 0;
		int max_model = -1;
		for (size_t j=0; j<models.size(); j++) {
			double prob = models[j]->viterbi(test[i].data_);
			if (prob > max_prob) {
				max_prob = prob;
				max_model = j;
			}
		}

		double garbage_prob = garbage_model->viterbi(test[i].data_);

		int test_type_id = types_s[test[i].type_];

		if (garbage_prob > max_prob) {
			without++;
			result[test_type_id][types.size()]++;
		}
		else if (max_model == test_type_id){
			success++;
			result[test_type_id][test_type_id]++;
		}
		else {
			result[test_type_id][max_model]++;
		}
	}
	cout << test.size() << " " << success << " " << without << endl;
	without_prob = 1.*without/test.size();

	delete garbage_model;
	for (size_t i=0; i<models.size(); i++) {
		delete models[i];
	}
	return 1.0*success/test.size();
}
